#!/usr/bin/env python3

from ete4 import Tree
from ete4.treeview import SeqMotifFace, TreeStyle, add_face_to_node


seq = ('-----------------------------------------------AQAK---IKGSKKAIKVFSSA---'
       'APERLQEYGSIFTDA---GLQRRPRHRIQSK-------ALQEKLKDFPVCVSTKPEPEDDAEEGLGGLPSN'
       'ISSVSSLLLFNTTENLYKKYVFLDPLAG----THVMLGAETEEKLFDAPLSISKREQLEQQVPENYFYVPD'
       'LGQVPEIDVPSYLPDLPGIANDLMYIADLGPGIAPSAPGTIPELPTFHTEVAEPLKVGELGSGMGAGPGTP'
       'AHTPSSLDTPHFVFQTYKMGAPPLPPSTAAPVGQGARQDDSSSSASPSVQGAPREVVDPSGGWATLLESIR'
       'QAGGIGKAKLRSMKERKLEKQQQKEQEQVRATSQGGHL--MSDLFNKLVMRRKGISGKGPGAGDGPGGAFA'
       'RVSDSIPPLPPPQQPQAEDED----')

mixed_motifs = [
        # seq.start, seq.end, shape, width, height, fgcolor, bgcolor
        [10, 100, '[]', None, 10, 'black', 'rgradient:blue', 'arial|8|white|long text clipped long text clipped'],
        [101, 150, 'o', None, 10, 'blue', 'pink', None],
        [155, 180, '()', None, 10, 'blue', 'rgradient:purple', None],
        [160, 190, '^', None, 14, 'black', 'yellow', None],
        [191, 200, '<>', None, 12, 'black', 'rgradient:orange', None],
        [201, 250, 'o', None, 12, 'black', 'brown', None],
        [351, 370, 'v', None, 15, 'black', 'rgradient:gold', None],
        [370, 420, 'compactseq', 2, 10, None, None, None],
]

simple_motifs = [
        # seq.start, seq.end, shape, width, height, fgcolor, bgcolor
        [10, 60, '[]', None, 10, 'black', 'rgradient:blue', 'arial|8|white|long text clipped long text clipped'],
        [120, 150, 'o', None, 10, 'blue', 'pink', None],
        [200, 300, '()', None, 10, 'blue', 'red', 'arial|8|white|hello'],
]

box_motifs = [
        # seq.start, seq.end, shape, width, height, fgcolor, bgcolor
        [0,  5, '[]', None, 10, 'black', 'rgradient:blue', 'arial|8|white|10'],
        [10, 25, '[]', None, 10, 'black', 'rgradient:ref', 'arial|8|white|10'],
        [30, 45, '[]', None, 10, 'black', 'rgradient:orange', 'arial|8|white|20'],
        [50, 65, '[]', None, 10, 'black', 'rgradient:pink', 'arial|8|white|20'],
        [70, 85, '[]', None, 10, 'black', 'rgradient:green', 'arial|8|white|20'],
        [90, 105, '[]', None, 10, 'black', 'rgradient:brown', 'arial|8|white|20'],
        [110, 125, '[]', None, 10, 'black', 'rgradient:yellow', 'arial|8|white|20'],
]


def get_example_tree():
        # Create a random tree and add to each leaf a random set of motifs
        # from the original set
        t = Tree('( (A, B, C, D, E, F, G), H, I);')

        seqFace = SeqMotifFace(seq, gapcolor='red')
        t['A'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq, seq_format='line', gap_format='blank')
        t['B'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq, seq_format='line')
        t['C'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq, seq_format='()')
        t['D'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq, motifs=simple_motifs, seq_format='-')
        t['E'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq=None, motifs=simple_motifs, gap_format='blank')
        t['F'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq, motifs=mixed_motifs, seq_format='-')
        t['G'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq=None, motifs=box_motifs, gap_format='line')
        t['H'].add_face(seqFace, 0, 'aligned')

        seqFace = SeqMotifFace(seq[30:60], seq_format='seq')
        t['I'].add_face(seqFace, 0, 'aligned')

        return t


if __name__ == '__main__':
    t = get_example_tree()
    ts = TreeStyle()
    ts.tree_width = 50
    t.render('seq_motif_faces.png', tree_style=ts)
